package algorithm;

import dataset.Dataset;
import exception.KNoMeaningException;
import scicalc.Calculate;
import scicalc.Collection;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class AGNES extends AbstractClusterAlgorithm {
    private int k = 3;
    private Linkage linkage = Linkage.AVERAGE_LINKAGE;

    @Override
    public long fit(Dataset X) {
        long start = System.currentTimeMillis();
        Map<Integer, List<Integer>> clusters = new HashMap<>();
        for (int i = 0; i < X.shape[0]; i++) {
            List<Integer> list = new ArrayList<>();
            list.add(i);
            clusters.put(i, list);
        }
        double[][] M = new double[X.shape[0]][X.shape[0]];
        for (int i = 0; i < X.shape[0]; i++) {
            M[i][i] = Double.MAX_VALUE;
            for (int j = i + 1; j < X.shape[0]; j++) {
                M[i][j] = linkage.distance(X.dataset, clusters.get(i), clusters.get(j));
                M[j][i] = M[i][j];
            }
        }
        int q = X.shape[0];
        while (q > k) {
            int index = Calculate.argmin(M);
            int i_ = index / M.length;
            int j_ = index % M.length;
            clusters.put(i_, Collection.union(clusters.get(i_), clusters.get(j_)));
            for (int i = j_ + 1; i < q; i++) {
                clusters.put(i - 1, clusters.get(i));
            }
            clusters.remove(q - 1);
            M = Calculate.delete(M, j_, Calculate.AXIS_0);
            M = Calculate.delete(M, j_, Calculate.AXIS_1);
            q--;
            for (int j = 0; j < q; j++) {
                if (i_ != j) {
                    M[i_][j] = linkage.distance(X.dataset, clusters.get(i_), clusters.get(j));
                    M[j][i_] = M[i_][j];
                }
            }
        }
        this.labels = new int[X.shape[0]];
        for (int i = 0; i < k; i++) {
            for (Integer j : clusters.get(i)) {
                this.labels[j] = i;
            }
        }
        this.result.shape[0] = X.shape[0];
        this.result.shape[1] = X.shape[1];
        this.result.dataset = X.dataset;
        this.result.label = labels;
        return System.currentTimeMillis() - start;
    }

    public void setK(int k) {
        if (k < 2) {
            try {
                throw new KNoMeaningException(k);
            } catch (KNoMeaningException e) {
                e.printStackTrace();
            }
        } else this.k = k;
    }

    public void setLinkage(Linkage linkage) {
        this.linkage = linkage;
    }
}
